#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (c) 2024 Huawei Device Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import datetime
import getpass
import os
import platform
import signal
import subprocess
import sys
import tarfile

DEFAULT_TIMEOUT = 60000
DEFAULT_PERMISSION = 'user'
DEFAULT_SWIFT = os.path.join(os.getcwd(), 'swift-5.9.2-RELEASE-ubuntu18.04', 'usr', 'bin', 'swift')
DEFAULT_JOB_COUNT = 36
DEFAULT_PROFILE = 'ark'
DEFAULT_STORAGE_PATH = 'storage2'

ARK_COMPILER_DIR = 'arkcompiler'
ETS_FRONTEND_DIR = 'ets_frontend'
ETS_RUNTIME_DIR = 'ets_runtime'
RUNTIME_CORE_DIR = 'runtime_core'
TOOL_CHAIN_DIR = 'toolchain'
ARK_FUZZILLI_DIR = 'ark_fuzzilli'
FUZZILLI_DIR = 'fuzzilli'
FUZZILLI_PATCH_DIR = 'patch'

FUZZILLI_GIT_HASH = '5696921ce2b021f5dacd56ba29f2a022c85b5891'
FUZZILLI_GIT_URL = 'https://gitee.com/mirrors_googleprojectzero/fuzzilli.git'
SWIFT_DOWNLOAD_URL = 'https://download.swift.org/swift-5.9.2-release/ubuntu1804/'\
                      'swift-5.9.2-RELEASE/swift-5.9.2-RELEASE-ubuntu18.04.tar.gz'


TERM_NORMAL = '\033[0m'
TERM_BLUE = '\033[1;34m'


def current_time():
    return datetime.datetime.now().strftime('%m-%d %H:%M:%S.%f')


class Logging():
    def __init__(self):
        self.is_logging = True

    def debug(self, info):
        if self.is_logging:
            print(f'{current_time()} D:>>>  {TERM_BLUE}{str(info)}{TERM_NORMAL}')

    def info(self, info):
        if self.is_logging:
            if len(str(info)) > 100:
                print(f'{current_time()} I:>>> \n{str(info)} ')
            else:
                print(f'{current_time()} I:>>>    {str(info)} ')


LOGGING = Logging()


class CommandCwd():
    def __init__(self, cmds, cwd):
        self.cmds = cmds
        self.cwd = cwd

    def run(self):
        cmd = ' '.join(self.cmds)
        LOGGING.debug('command: ' + cmd + ' | ' + 'dir: ' + self.cwd)
        if platform.system() == 'Windows' :
            proc = subprocess.Popen(self.cmds, cwd=self.cwd ,shell=True)
        else :
            proc = subprocess.Popen(self.cmds, cwd=self.cwd)
        return proc.wait()


def run_cmd_cwd(commands, cwd=os.getcwd()):
    cmd = CommandCwd(commands, cwd)
    return cmd.run()


def git_clone(git_url, code_dir):
    cmd = ['git', 'clone', git_url, code_dir]
    ret = run_cmd_cwd(cmd)
    assert not ret, f"\n error: Cloning '{git_url}' failed."


def git_checkout(git_bash, cwd):
    cmd = ['git', 'checkout', git_bash]
    ret = run_cmd_cwd(cmd, cwd)
    assert not ret, f"\n error: git checkout '{git_bash}' failed."


def git_apply(patch_file, cwd):
    cmd = ['git', 'apply', patch_file]
    ret = run_cmd_cwd(cmd, cwd)
    assert not ret, f"\n error: Failed to apply '{patch_file}'"


def get_dir_abs_path(dir):
    if not os.path.isdir(dir):
        sys.stderr.write(f"The directory {dir} does not exist.\n")
        return ''

    return os.path.abspath(dir)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--swift', help='the swift for fuzzilli test')
    parser.add_argument(
        '--jobs', default=DEFAULT_JOB_COUNT, help='the jobs count for fuzzilli test')
    parser.add_argument(
        '--profile', default=DEFAULT_PROFILE, help='the profile for fuzzilli test')
    parser.add_argument(
        '--storagePath', default=DEFAULT_STORAGE_PATH, help='the storage path for fuzzilli test')

    return parser.parse_args()


def exec_command(cmd_args, timeout=DEFAULT_TIMEOUT, permission=DEFAULT_PERMISSION):
    msg = ''
    retcode = 0
    is_sudo = permission == 'sudo'

    if is_sudo:
        password = getpass.getpass('Enter your sudo password: ')
        cmd_args = ['sudo', '-S'] + cmd_args

    proc = subprocess.Popen(cmd_args,
                            stderr=subprocess.PIPE,
                            stdout=subprocess.PIPE,
                            stdin=subprocess.PIPE if is_sudo else None,
                            close_fds=True,
                            start_new_session=True,
                            text=True)

    cmd_string = ' '.join(cmd_args)
    print(f"Executing command: {cmd_string}")

    try:
        if is_sudo:
            std_out, std_err = proc.communicate(password + '\n', timeout=timeout)
        else:
            std_out, std_err = proc.communicate(timeout=timeout)
        retcode = proc.poll()

        if std_err:
            sys.stderr.write(f"Error: {std_err}.\n")

        if std_out:
            print(f"Output: {std_out}")

    except subprocess.TimeoutExpired:
        proc.kill()
        proc.terminate()
        os.kill(proc.pid, signal.SIGTERM)
        retcode = 1
        msg = f"Timeout: '{cmd_string}' timed out after {timeout} seconds.\n"
    except Exception as err:
        retcode = 1
        msg = f"{cmd_string}: unknown error: {str(err)}.\n"
    finally:
        if proc.stdout:
            proc.stdout.close()
        if proc.stderr:
            proc.stderr.close()
        proc.wait()

    if msg:
        sys.stderr.write(msg)

    return retcode, std_out


class PrepareSwiftCode():
    def check_swift_version(self, swift_bin_path):
        cmd = [swift_bin_path, '--version']
        retcode, std_out = exec_command(cmd)
        return retcode


    def prepare_swift_code(self, retcode):
        print('Downloading swift...')
        current_dir = os.getcwd()
        swift_path = os.path.join(current_dir, 'swift.tar.gz')
        cmd = ['wget', '-O', swift_path, SWIFT_DOWNLOAD_URL]
        retcode, std_out = exec_command(cmd)
        if retcode:
            sys.stderr.write('Failed to download swift.\n')
            return retcode
        print('Downloading swift finished.')

        with tarfile.open(swift_path, 'r:gz') as tar:
            print('Extract swift.tar.gz files...')
            tar.extractall(current_dir)

        return retcode


    def run(self, retcode, args):
        if args.swift:
            return self.check_swift_version(args.swift)

        return self.prepare_swift_code(retcode)


class PrepareFuzzilliCode():
    def apply_fuzzilli_patch(self):
        current_dir = os.getcwd()
        fuzzilli_patch_path = os.path.join(current_dir, ARK_FUZZILLI_DIR, FUZZILLI_PATCH_DIR, 'fuzzilli.diff')

        if os.path.exists(fuzzilli_patch_path) and os.path.isfile(fuzzilli_patch_path):
            git_apply(fuzzilli_patch_path, FUZZILLI_DIR)


    def prepare_fuzzilli_code(self, retcode):
        if os.path.isdir(os.path.join(FUZZILLI_DIR, '.git')):
            print('Computer has fuzzilli code.')
            return retcode

        git_clone(FUZZILLI_GIT_URL, FUZZILLI_DIR)
        git_checkout(FUZZILLI_GIT_HASH, FUZZILLI_DIR)
        self.apply_fuzzilli_patch()

        return retcode


    def get_sysctl_config(self):
        cmd = ['sysctl', '-a', '|grep', 'kernel.core_pattern']
        retcode, std_out = exec_command(cmd, DEFAULT_TIMEOUT, 'sudo')
        if retcode:
            sys.stderr.write('Failed to get sysctl config.\n')
            return ''

        lines = std_out.split('\n')
        config = {}
        for line in lines:
            if line.strip() != '':
                key, value = line.split(' = ')
                config[key.strip()] = value.strip()

        return config


    def check_bin_false(self):
        config = self.get_sysctl_config()
        if not config:
            sys.stderr.write('Failed to get sysctl config.\n')
            return False

        kcp_val = config.get('kernel.core_pattern')
        return kcp_val == '|/bin/false'


    def init_kernel_core_pattern(self, retcode):
        is_bin_false = self.check_bin_false()
        if is_bin_false:
            print('The kernel.core_pattern is bin/false')
            return retcode

        cmd = ['sysctl', '-w', 'kernel.core_pattern=|/bin/false']
        retcode, std_out = exec_command(cmd, DEFAULT_TIMEOUT, 'sudo')
        return retcode


    def run(self, retcode):
        if self.prepare_fuzzilli_code(retcode):
            return retcode

        if self.init_kernel_core_pattern(retcode):
            sys.stderr.write('Failed to initialize kernel-core pattern.\n')
            return retcode

        return retcode


class FuzzilliTest():
    def apply_arkcompiler_patch(self):
        current_dir = os.getcwd()
        patches = [
            ('ets_frontend.diff', ETS_FRONTEND_DIR),
            ('ets_runtime.diff', ETS_RUNTIME_DIR),
            ('runtime_core.diff', RUNTIME_CORE_DIR),
            ('toolchain.diff', TOOL_CHAIN_DIR),
        ]

        for patch_file, target_dir in patches:
            patch_path = os.path.join(current_dir, ARK_FUZZILLI_DIR, FUZZILLI_PATCH_DIR, patch_file)
            if os.path.exists(patch_path) and os.path.isfile(patch_path):
                git_apply(patch_path, os.path.join(current_dir, ARK_COMPILER_DIR, target_dir))


    def run_fuzzilli_test(self, retcode, args):
        build_standalone_cmd = ['python', 'ark.py', 'x64.release', 'arkfuzzer']
        retcode, std_out = exec_command(build_standalone_cmd)
        if retcode:
            return retcode

        profile = args.profile if args.profile else DEFAULT_PROFILE
        jobs = args.jobs if args.jobs else DEFAULT_JOB_COUNT
        storage = args.storagePath if args.storagePath else DEFAULT_STORAGE_PATH
        swift_tool = args.swift if args.swift else DEFAULT_SWIFT
        current_dir = os.getcwd()
        output_dir = os.path.join(current_dir, FUZZILLI_DIR, storage)
        arkfuzzili_tool = os.path.join(current_dir, 'out', 'x64.release', 'arkfuzzer', 'arkfuzzer', 'arkfuzzer')

        os.chdir(os.path.join(current_dir, FUZZILLI_DIR))
        run_fuzzilli_cmd = [swift_tool, 'run', 'FuzzilliCli', f'--profile={profile}', f'--jobs={jobs}', '--overwrite',
                            f'--storagePath={output_dir}', arkfuzzili_tool]
        retcode, std_out = exec_command(run_fuzzilli_cmd)

        return retcode


    def run(self, retcode, args):
        self.apply_arkcompiler_patch()
        if self.run_fuzzilli_test(retcode, args):
            return retcode

        return retcode


def prepare_swift(retcode, args):
    swift_prepare = PrepareSwiftCode()
    swift_prepare.run(retcode, args)

    return retcode


def prepare_fuzzilli(retcode):
    fuzzilli_prepare = PrepareFuzzilliCode()
    fuzzilli_prepare.run(retcode)

    return retcode


def run_fuzzilli(retcode, args):
    fuzzilli_test = FuzzilliTest()
    fuzzilli_test.run(retcode, args)

    return retcode


def main(args):
    retcode = 0
    if prepare_swift(retcode, args):
        sys.stderr.write('Failed to prepare swift code.\n')
        return retcode

    if prepare_fuzzilli(retcode):
        sys.stderr.write('Failed to prepare fuzzilli code.\n')
        return retcode

    if run_fuzzilli(retcode, args):
        sys.stderr.write('Failed to run fuzzilli test.\n')
        return retcode

    return retcode


if __name__ == '__main__':
    main(parse_args())
